#!/usr/bin/env python3
"""
WebSocket real-time operations example.

This example demonstrates how to use WebSocket connections for real-time
vector operations, including live search, embedding generation, and
monitoring.
"""

import asyncio
import json
import websockets
import time
from typing import List, Dict, Any


class VectorDBWebSocketClient:
    """WebSocket client for VectorDB real-time operations."""

    def __init__(self, ws_url: str = "ws://localhost:8000/ws"):
        self.ws_url = ws_url
        self.websocket = None
        self.connection_id = None
        self.message_handlers = {}
        self.running = False

    async def connect(self):
        """Connect to WebSocket server."""
        print(f"🔌 Connecting to {self.ws_url}...")
        self.websocket = await websockets.connect(self.ws_url)
        self.running = True

        # Wait for welcome message
        welcome_msg = await self.receive_message()
        if welcome_msg and welcome_msg.get("type") == "welcome":
            self.connection_id = welcome_msg["data"]["connection_id"]
            print(f"✅ Connected! Connection ID: {self.connection_id}")
        else:
            raise Exception("Failed to receive welcome message")

    async def disconnect(self):
        """Disconnect from WebSocket server."""
        self.running = False
        if self.websocket:
            await self.websocket.close()
            print("🔌 Disconnected from WebSocket server")

    async def send_message(self, message_type: str, data: Dict[str, Any], request_id: str = None):
        """Send message to server."""
        if not self.websocket:
            raise Exception("Not connected to WebSocket server")

        message = {
            "type": message_type,
            "data": data,
            "request_id": request_id or f"req_{int(time.time() * 1000)}"
        }

        await self.websocket.send(json.dumps(message))
        return message["request_id"]

    async def receive_message(self) -> Dict[str, Any]:
        """Receive message from server."""
        if not self.websocket:
            raise Exception("Not connected to WebSocket server")

        try:
            message_text = await self.websocket.recv()
            return json.loads(message_text)
        except websockets.exceptions.ConnectionClosed:
            self.running = False
            return None

    async def ping(self) -> float:
        """Send ping and measure response time."""
        start_time = time.time()
        request_id = await self.send_message("ping", {})

        # Wait for pong response
        while self.running:
            message = await self.receive_message()
            if message and message.get("type") == "pong" and message.get("request_id") == request_id:
                return time.time() - start_time

        return -1.0

    async def health_check(self) -> Dict[str, Any]:
        """Perform health check."""
        request_id = await self.send_message("health", {})

        # Wait for health response
        while self.running:
            message = await self.receive_message()
            if message and message.get("type") == "health_response" and message.get("request_id") == request_id:
                return message["data"]

        return {}

    async def embed_texts(self, texts: List[str], provider: str = None) -> List[List[float]]:
        """Generate embeddings for texts."""
        request_id = await self.send_message("embed", {
            "texts": texts,
            "provider": provider
        })

        # Wait for embed response
        while self.running:
            message = await self.receive_message()
            if message and message.get("type") == "embed_response" and message.get("request_id") == request_id:
                return message["data"]["embeddings"]

        return []

    async def search_vectors(self, table_name: str, query: str,
                           limit: int = 10, threshold: float = 0.0,
                           provider: str = None) -> List[Dict[str, Any]]:
        """Search for similar vectors."""
        request_id = await self.send_message("search", {
            "table_name": table_name,
            "query": query,
            "limit": limit,
            "threshold": threshold,
            "provider": provider
        })

        # Wait for search response
        while self.running:
            message = await self.receive_message()
            if message and message.get("type") == "search_response" and message.get("request_id") == request_id:
                return message["data"]["results"]

        return []

    async def store_vectors(self, table_name: str, texts: List[str],
                          metadata: List[Dict[str, Any]] = None,
                          provider: str = None) -> bool:
        """Store vectors in database."""
        request_id = await self.send_message("store", {
            "table_name": table_name,
            "texts": texts,
            "metadata": metadata or [],
            "provider": provider
        })

        # Wait for store response
        while self.running:
            message = await self.receive_message()
            if message and message.get("type") == "store_response" and message.get("request_id") == request_id:
                return message["data"]["success"]

        return False

    async def get_stats(self) -> Dict[str, Any]:
        """Get connection statistics."""
        request_id = await self.send_message("stats", {})

        # Wait for stats response
        while self.running:
            message = await self.receive_message()
            if message and message.get("type") == "stats_response" and message.get("request_id") == request_id:
                return message["data"]

        return {}

    async def subscribe_to_events(self, event_types: List[str]):
        """Subscribe to event types."""
        for event_type in event_types:
            await self.send_message("subscribe", {"type": event_type})

    async def unsubscribe_from_events(self, event_types: List[str]):
        """Unsubscribe from event types."""
        for event_type in event_types:
            await self.send_message("unsubscribe", {"type": event_type})

    async def listen_for_events(self, event_handler: callable):
        """Listen for events and handle them."""
        while self.running:
            try:
                message = await self.receive_message()
                if message and message.get("type", "").startswith("event_"):
                    event_type = message["type"][6:]  # Remove "event_" prefix
                    await event_handler(event_type, message.get("data", {}))
            except Exception as e:
                print(f"Error handling event: {e}")
                break


async def basic_websocket_example():
    """Basic WebSocket operations example."""
    print("🔌 WebSocket Basic Operations Example")
    print("=" * 50)

    client = VectorDBWebSocketClient()

    try:
        # Connect to WebSocket
        await client.connect()

        # 1. Health check
        print("1. Performing health check...")
        health = await client.health_check()
        print(f"   Service status: {health.get('status', 'unknown')}")
        print(f"   Database healthy: {health.get('database', {}).get('healthy', False)}")
        print(f"   Embedding healthy: {health.get('embedding', {}).get('healthy', False)}")
        print()

        # 2. Ping test
        print("2. Testing connection latency...")
        ping_time = await client.ping()
        print(f"   Ping response time: {ping_time:.3f} seconds")
        print()

        # 3. Embed texts
        print("3. Generating embeddings via WebSocket...")
        texts = ["Hello world", "WebSocket is fast", "Real-time operations"]
        embeddings = await client.embed_texts(texts)

        print(f"   Generated {len(embeddings)} embeddings")
        print(f"   Embedding dimensions: {len(embeddings[0]) if embeddings else 0}")
        print(f"   First embedding (first 5 values): {embeddings[0][:5] if embeddings else 'N/A'}")
        print()

        # 4. Create table and store vectors
        print("4. Creating table and storing vectors...")
        await client.send_message("store", {
            "table_name": "websocket_test",
            "texts": ["WebSocket test document 1", "WebSocket test document 2"],
            "metadata": [{"source": "websocket"}, {"source": "websocket"}]
        })

        # Wait for store response
        while client.running:
            message = await client.receive_message()
            if message and message.get("type") == "store_response":
                success = message["data"]["success"]
                print(f"   Store operation: {'Success' if success else 'Failed'}")
                break
        print()

        # 5. Search vectors
        print("5. Searching vectors via WebSocket...")
        results = await client.search_vectors("websocket_test", "websocket", limit=5)

        print(f"   Found {len(results)} results")
        for i, result in enumerate(results, 1):
            print(f"   {i}. {result['text']}")
            print(f"      Score: {result['score']:.4f}")
        print()

        # 6. Get connection stats
        print("6. Connection statistics...")
        stats = await client.get_stats()
        print(f"   Total connections: {stats.get('total_connections', 0)}")
        print(f"   This connection messages: {stats.get('connections', {}).get(client.connection_id, {}).get('message_count', 0)}")
        print()

        print("✅ Basic WebSocket example completed successfully!")

    except Exception as e:
        print(f"❌ Error: {e}")
        raise

    finally:
        await client.disconnect()


async def real_time_monitoring_example():
    """Real-time monitoring and event handling example."""
    print("📊 Real-time Monitoring Example")
    print("=" * 50)

    client = VectorDBWebSocketClient()

    # Event handler
    async def handle_event(event_type: str, data: Dict[str, Any]):
        timestamp = data.get('timestamp', 'N/A')
        print(f"📡 Event: {event_type} at {timestamp}")

        if event_type == "new_vectors":
            print(f"   New vectors stored: {data.get('count', 0)} in table '{data.get('table_name', 'unknown')}'")
        elif event_type == "search_completed":
            print(f"   Search completed: '{data.get('query', 'unknown')}' returned {data.get('result_count', 0)} results")
        elif event_type == "table_created":
            print(f"   New table created: '{data.get('table_name', 'unknown')}'")
        elif event_type == "table_deleted":
            print(f"   Table deleted: '{data.get('table_name', 'unknown')}'")

    try:
        # Connect to WebSocket
        await client.connect()

        # Subscribe to events
        print("1. Subscribing to real-time events...")
        await client.subscribe_to_events([
            "new_vectors",
            "search_completed",
            "table_created",
            "table_deleted"
        ])
        print("   Subscribed to events: new_vectors, search_completed, table_created, table_deleted")
        print()

        # Start event listener in background
        event_task = asyncio.create_task(client.listen_for_events(handle_event))

        # Simulate activity
        print("2. Simulating database activity...")
        print()

        # Create a table
        print("   Creating test table...")
        await client.send_message("store", {
            "table_name": "monitoring_test",
            "texts": ["Monitoring test document"],
            "metadata": [{"test": True}]
        })
        await asyncio.sleep(1)

        # Store multiple documents
        print("   Storing multiple documents...")
        for i in range(3):
            await client.store_vectors("monitoring_test", [f"Document {i}"], [{"batch": i}])
            await asyncio.sleep(0.5)

        # Perform searches
        print("   Performing searches...")
        for query in ["test", "document", "monitoring"]:
            results = await client.search_vectors("monitoring_test", query, limit=2)
            await asyncio.sleep(0.5)

        # Wait a bit for events to process
        await asyncio.sleep(2)

        # Unsubscribe from events
        print("3. Unsubscribing from events...")
        await client.unsubscribe_from_events(["table_deleted"])  # Unsubscribe from one event
        print()

        # Stop event listener
        event_task.cancel()
        try:
            await event_task
        except asyncio.CancelledError:
            pass

        print("✅ Real-time monitoring example completed successfully!")

    except Exception as e:
        print(f"❌ Error: {e}")
        raise

    finally:
        await client.disconnect()


async def concurrent_operations_example():
    """Concurrent WebSocket operations example."""
    print("⚡ Concurrent Operations Example")
    print("=" * 50)

    # Create multiple clients
    clients = [VectorDBWebSocketClient() for _ in range(3)]

    try:
        # Connect all clients
        print("1. Connecting multiple clients...")
        connect_tasks = [client.connect() for client in clients]
        await asyncio.gather(*connect_tasks)
        print(f"   Connected {len(clients)} clients")
        print()

        # Prepare test data
        test_data = [
            ("concurrent_test", ["Concurrent test 1", "Concurrent test 2"]),
            ("concurrent_test", ["Concurrent test 3", "Concurrent test 4"]),
            ("concurrent_test", ["Concurrent test 5", "Concurrent test 6"])
        ]

        # Concurrent store operations
        print("2. Concurrent store operations...")
        store_tasks = []
        for i, (table_name, texts) in enumerate(test_data):
            task = clients[i].store_vectors(table_name, texts, [{"client": i} for _ in texts])
            store_tasks.append(task)

        store_results = await asyncio.gather(*store_tasks)
        success_count = sum(1 for result in store_results if result)
        print(f"   Store operations completed: {success_count}/{len(store_tasks)} successful")
        print()

        # Concurrent search operations
        print("3. Concurrent search operations...")
        search_queries = ["concurrent", "test", "operations"]
        search_tasks = []

        for i, query in enumerate(search_queries):
            task = clients[i].search_vectors("concurrent_test", query, limit=3)
            search_tasks.append(task)

        search_results = await asyncio.gather(*search_tasks)

        for i, results in enumerate(search_results):
            print(f"   Client {i} search results for '{search_queries[i]}': {len(results)} results")
        print()

        # Concurrent embedding operations
        print("4. Concurrent embedding operations...")
        embed_texts = [
            ["Client 1 embedding text"],
            ["Client 2 embedding text", "Another text for client 2"],
            ["Client 3 text 1", "Client 3 text 2", "Client 3 text 3"]
        ]

        embed_tasks = []
        for i, texts in enumerate(embed_texts):
            task = clients[i].embed_texts(texts)
            embed_tasks.append(task)

        embed_results = await asyncio.gather(*embed_tasks)

        for i, embeddings in enumerate(embed_results):
            print(f"   Client {i} generated {len(embeddings)} embeddings")
        print()

        # Concurrent health checks
        print("5. Concurrent health checks...")
        health_tasks = [client.health_check() for client in clients]
        health_results = await asyncio.gather(*health_tasks)

        for i, health in enumerate(health_results):
            print(f"   Client {i} health: {health.get('status', 'unknown')}")
        print()

        print("✅ Concurrent operations example completed successfully!")

    except Exception as e:
        print(f"❌ Error: {e}")
        raise

    finally:
        # Disconnect all clients
        disconnect_tasks = [client.disconnect() for client in clients]
        await asyncio.gather(*disconnect_tasks, return_exceptions=True)


async def performance_benchmark_example():
    """Performance benchmarking example."""
    print("🏃 Performance Benchmark Example")
    print("=" * 50)

    client = VectorDBWebSocketClient()

    try:
        # Connect to WebSocket
        await client.connect()

        # Prepare test data
        print("1. Preparing test data...")
        test_texts = [f"Performance test document {i}" for i in range(100)]
        test_metadata = [{"test_id": i, "benchmark": True} for i in range(100)]

        # Store test data
        print("2. Storing test data...")
        start_time = time.time()
        success = await client.store_vectors("benchmark_test", test_texts, test_metadata)
        store_time = time.time() - start_time
        print(f"   Store 100 documents: {store_time:.3f}s")
        print()

        # Benchmark operations
        print("3. Benchmarking operations...")
        operations = [
            ("ping", lambda: client.ping()),
            ("health_check", lambda: client.health_check()),
            ("embed_single", lambda: client.embed_texts(["benchmark text"])),
            ("search", lambda: client.search_vectors("benchmark_test", "performance", limit=10)),
            ("stats", lambda: client.get_stats())
        ]

        results = {}
        for op_name, op_func in operations:
            # Warm up
            await op_func()

            # Benchmark
            times = []
            for _ in range(5):  # 5 iterations
                start_time = time.time()
                await op_func()
                times.append(time.time() - start_time)

            avg_time = sum(times) / len(times)
            results[op_name] = {
                "avg_time": avg_time,
                "min_time": min(times),
                "max_time": max(times)
            }

            print(f"   {op_name}: avg={avg_time:.3f}s, min={min(times):.3f}s, max={max(times):.3f}s")

        print()

        # Summary
        print("4. Performance Summary...")
        print("   Operation Performance (average time):")
        for op_name, stats in results.items():
            print(f"   - {op_name}: {stats['avg_time']:.3f}s")
        print()

        print("✅ Performance benchmark completed successfully!")

    except Exception as e:
        print(f"❌ Error: {e}")
        raise

    finally:
        await client.disconnect()


async def main():
    """Run all WebSocket examples."""
    print("VectorDB WebSocket Examples")
    print("=" * 60)
    print()

    # Run basic WebSocket example
    await basic_websocket_example()
    print()
    print("-" * 60)
    print()

    # Run real-time monitoring example
    await real_time_monitoring_example()
    print()
    print("-" * 60)
    print()

    # Run concurrent operations example
    await concurrent_operations_example()
    print()
    print("-" * 60)
    print()

    # Run performance benchmark example
    await performance_benchmark_example()


if __name__ == "__main__":
    asyncio.run(main())